#include <stdio.h>
#include <arm_neon.h>
#include <stdlib.h>
#include <sys/time.h>
#include <stdbool.h>
void print_matrix(double *data,int m,int n){
    for(int i=0;i<m;i++){
        for(int j=0;j<n;j++){
            printf("%f ",data[i*n+j]);
        }
        printf("\n");
    }
}
bool is_matrix_same(float32_t *a,float32_t *b, int size){
    for(int i=0;i<size;i++){
        if(a[i]!=b[i]){
            printf("matrices are difference at %d\n",i);
            return false;
        }
    }
    return true;
}

double what_time_is_it_now()
{
    struct timeval time;
    if (gettimeofday(&time,NULL)){
        return 0;
    }
    return (double)time.tv_sec + (double)time.tv_usec * .000001;
}
void rand_init_matrix(float32_t *m,int size){
    for(int i=0;i<size;i++){
        m[i]=(float)rand()/RAND_MAX;
    }
}

//这个方法是正常的C语言实现
double *fvec_mul_fmatrix(float *rand_vec, float *a, uint32_t M, uint32_t K)
{
    //double *result = TEE_Malloc(sizeof(double) * K, 0);
    double *result = (double*)calloc(K,sizeof(double));
    if (!result){
        //EMSG("Malloc ff_rec of %fKB size error.",sizeof(double)*K/1024.0);
        return NULL;
    }  
    ////TEE_MemFill(result, 0, sizeof(double) * K);

    uint32_t tmp = 0;
    // Caculate res_ra=rand_vec*A
    for (uint32_t row = 0; row < M; row++)
    {
        //tmp = row * K;
        for (uint32_t col = 0; col < K; col++)
        {
            result[col] += rand_vec[row] * (double)a[tmp];
            tmp++;
        }
    }
    return result;
}
//这个方法是以行来加载矩阵a的数据到寄存器
//但是这个方法就需要多次将结果返回到内存中的结果向量中
double *fvec_mul_fmatrix_neon1(float32_t *vec, float32_t *a, uint32_t M, uint32_t K){
    //double *result = TEE_Malloc(sizeof(double) * K, 0);
    double *result = (double*)calloc(K,sizeof(double));
    if (!result){
        //EMSG("Malloc ff_rec of %fKB size error.",sizeof(double)*K/1024.0);
        return NULL;
    } 
    uint32_t Kn=K>>1;
    float32_t *atmp;
    float32x2_t tmp;
    float64x2_t c0,a0;
    for(uint32_t x=0;x<M;++x){
        float64_t vtmp=vec[x];
        atmp=a+x*K;
        for(uint32_t y=0;y<Kn;++y){
            tmp=vld1_f32(atmp+(y<<1));
            c0=vld1q_f64(result+(y<<1));
            a0=vcvt_f64_f32(tmp);
            c0=vaddq_f64(c0,vmulq_n_f64(a0,vtmp));
            vst1q_f64(result+(y<<1),c0);
        }
        if(K&1) result[K-1]+=vtmp*atmp[K-1];
    }
    return result;
}

double *fvec_mul_fmatrix_neon2(float32_t *vec, float32_t *a, uint32_t M, uint32_t K){
    //double *result = TEE_Malloc(sizeof(double) * K, 0);
    double *result = (double*)calloc(K,sizeof(double));
    if (!result){
        //EMSG("Malloc ff_rec of %fKB size error.",sizeof(double)*K/1024.0);
        return NULL;
    } 
    uint32_t mn=M-1;
    uint32_t kn=K-1;
    uint32_t flag1=M&1,flag2=K&1;
    float32_t *ap;
    float32x2_t tmp0,tmp1,tmp2;
    float64x2_t c0,a0,a1,v0;
    for(uint32_t x=0;x<kn;x+=2){
        ap=a+x;
        c0=vmovq_n_f64(0);
        for(uint32_t y=0,idy=0;y<mn;y+=2,idy=y*K){
            tmp0=vld1_f32(vec+y);
            tmp1=vld1_f32(ap+idy);
            tmp2=vld1_f32(ap+idy+K);
            v0=vcvt_f64_f32(tmp0);
            a0=vcvt_f64_f32(tmp1);
            a1=vcvt_f64_f32(tmp2);
            c0=vfmaq_laneq_f64(c0,a0,v0,0);
            c0=vfmaq_laneq_f64(c0,a1,v0,1);
        }
        if(flag1){
            float32x2_t tmp3=vld1_f32(ap+mn*K);
            float64x2_t a0=vcvt_f64_f32(tmp3);
            c0=vaddq_f64(c0,vmulq_n_f64(a0,vec[mn]));
        }
        vst1q_f64(result+x,c0);
    }
    if(flag2){
        double tmpre=0;
        for(uint32_t x=0;x<M;x++){
            tmpre+=(double)vec[x]*a[kn+x*K];
        }
        result[kn]=tmpre;
    }

    return result;
}


bool is_matrix_same_double(double *a,double *b, int size){
    for(int i=0;i<size;i++){
        if(a[i]!=b[i]){
            printf("matrices are difference at %d\n",i);
            double tmp1=a[i];
            double tmp2=b[i];
            return false;
        }
    }
    printf("matrices are same\n");
    return true;
}



void test01(){
    uint32_t M=4003;
    uint32_t K=5001;
    float *vec=(float*)malloc(sizeof(float)*M);
    rand_init_matrix(vec,M);
    float *a=(float*)malloc(sizeof(float)*M*K);
    rand_init_matrix(a,M*K);
    //print_matrix()
    double beg1=what_time_is_it_now();
    double *result1=fvec_mul_fmatrix(vec,a,M,K);
    double end1=what_time_is_it_now();
    double *result2=fvec_mul_fmatrix_neon1(vec,a,M,K);
    double end2=what_time_is_it_now();
    double *result3=fvec_mul_fmatrix_neon2(vec,a,M,K);
    double end3=what_time_is_it_now();
    // print_matrix(result1,1,K);
    // printf("\n=======================================\n");
    // print_matrix(result2,1,K);
    // printf("\n=======================================\n");
    // print_matrix(result3,1,K);
    is_matrix_same_double(result1,result2,K);
    is_matrix_same_double(result1,result3,K);
    printf("c code time is %lf\n",end1-beg1);
    printf("neon1 code time is %lf\n",end2-end1);
    printf("neon2 code time is %lf\n",end3-end2);
}

int main(){
    test01();
}